{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A3C - Multi-agent\n",
    "\n",
    "While training is taking place, statistics on agent performance are available from Tensorboard. To launch it use:\n",
    "\n",
    "`tensorboard --logdir=train_RNN`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import random\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.layers as layers\n",
    "import matplotlib.pyplot as plt\n",
    "import threading\n",
    "import time\n",
    "import scipy.signal as signal\n",
    "import os\n",
    "import GroupLock\n",
    "import multiprocessing\n",
    "%matplotlib inline\n",
    "import minecraft\n",
    "import pickle\n",
    "\n",
    "from tensorflow.python.client import device_lib\n",
    "dev_list = device_lib.list_local_devices()\n",
    "print(dev_list)\n",
    "assert len(dev_list) > 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_gif(images, fname, duration=2, true_image=False,salience=False,salIMGS=None):\n",
    "    import moviepy.editor as mpy\n",
    "\n",
    "    def make_frame(t):\n",
    "        try:\n",
    "            x = images[int(len(images)/duration*t)]\n",
    "        except:\n",
    "            x = images[-1]\n",
    "\n",
    "        if true_image:\n",
    "            return x.astype(np.uint8)\n",
    "        else:\n",
    "            return ((x+1)/2*255).astype(np.uint8)\n",
    "\n",
    "    def make_mask(t):\n",
    "        try:\n",
    "            x = salIMGS[int(len(salIMGS)/duration*t)]\n",
    "        except:\n",
    "            x = salIMGS[-1]\n",
    "        return x\n",
    "\n",
    "    clip = mpy.VideoClip(make_frame, duration=duration)\n",
    "    if salience == True:\n",
    "        mask = mpy.VideoClip(make_mask, ismask=True,duration= duration)\n",
    "        clipB = clip.set_mask(mask)\n",
    "        clipB = clip.set_opacity(0)\n",
    "        mask = mask.set_opacity(0.1)\n",
    "        mask.write_gif(fname, fps = len(images) / duration,verbose=False)\n",
    "    else:\n",
    "        clip.write_gif(fname, fps = len(images) / duration,verbose=False)\n",
    "\n",
    "\n",
    "\n",
    "# Copies one set of variables to another.\n",
    "# Used to set worker network parameters to those of global network.\n",
    "def update_target_graph(from_scope,to_scope):\n",
    "    from_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, from_scope)\n",
    "    to_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, to_scope)\n",
    "\n",
    "    op_holder = []\n",
    "    for from_var,to_var in zip(from_vars,to_vars):\n",
    "        op_holder.append(to_var.assign(from_var))\n",
    "    return op_holder\n",
    "\n",
    "def discount(x, gamma):\n",
    "    return signal.lfilter([1], [1, -gamma], x[::-1], axis=0)[::-1]\n",
    "\n",
    "def good_discount(x, gamma):\n",
    "    return discount(x,gamma)\n",
    "\n",
    "#Used to initialize weights for policy and value output layers (Do we need to use that? Maybe not now)\n",
    "def normalized_columns_initializer(std=1.0):\n",
    "    def _initializer(shape, dtype=None, partition_info=None):\n",
    "        out = np.random.randn(*shape).astype(np.float32)\n",
    "        out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True))\n",
    "        return tf.constant(out)\n",
    "    return _initializer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A3C Approach"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing the Actor-Critic Network### Implementing the Deep Q-Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ACNet:\n",
    "    def __init__(self, scope, a_size, trainer):\n",
    "        \n",
    "        with tf.variable_scope(str(scope)+'/qvalues'):\n",
    "            self.is_Train = True\n",
    "            #The input size may require more work to fit the interface.\n",
    "            self.inputs = tf.placeholder(shape=[1,5,10,4,10], dtype=tf.float32)\n",
    "            self.myinput = tf.transpose(self.inputs, perm=[0,2,3,4,1])\n",
    "            self.policy, self.value, self.state_out, self.state_in, self.state_init, self.has_block, _ , self.is_built = self._build_net(self.myinput)\n",
    "        with tf.variable_scope(str(scope)+'/qvaluesB'):\n",
    "            self.inputsB = tf.placeholder(shape=[EXPERIENCE_BUFFER_SIZE,5,10,4,10], dtype=tf.float32)\n",
    "            self.myinputB = tf.transpose(self.inputsB,perm=[0,2,3,4,1])\n",
    "            self.policyB, self.valueB, self.state_outB, self.state_inB, self.state_initB, self.has_blockB, self.validsB, self.is_builtB = self._build_net(self.myinputB)\n",
    "        if(scope!=GLOBAL_NET_SCOPE):\n",
    "            self.actions = tf.placeholder(shape=[None], dtype=tf.int32)\n",
    "            self.actions_onehot = tf.one_hot(self.actions, a_size, dtype=tf.float32)\n",
    "            self.valids = tf.placeholder(shape=[None,a_size], dtype=tf.float32)\n",
    "            self.target_v = tf.placeholder(tf.float32, [None], 'Vtarget')\n",
    "            self.target_has = tf.placeholder(tf.float32, [None])\n",
    "            self.target_built = tf.placeholder(tf.float32, [None])\n",
    "            self.advantages = tf.placeholder(shape=[None], dtype=tf.float32)\n",
    "            self.responsible_outputs = tf.reduce_sum(self.policyB * self.actions_onehot, [1])\n",
    "            self.train_value = tf.placeholder(tf.float32, [None])\n",
    "\n",
    "            # Loss Functions\n",
    "            self.value_loss = 0.5 * tf.reduce_sum(self.train_value*tf.square(self.target_v - tf.reshape(self.valueB, shape=[-1])))\n",
    "\n",
    "            # something to encourage exploration\n",
    "            self.entropy = - tf.reduce_sum(self.policyB * tf.log(tf.clip_by_value(self.policyB,1e-10,1.0)))\n",
    "            self.block_loss = - tf.reduce_sum(self.target_has*tf.log(tf.clip_by_value(tf.reshape(self.has_blockB, shape=[-1]),1e-10,1.0))+(1-self.target_has)*tf.log(tf.clip_by_value(1-tf.reshape(self.has_blockB, shape=[-1]),1e-10,1.0)))\n",
    "            self.built_loss = - tf.reduce_sum(self.target_built*tf.log(tf.clip_by_value(tf.reshape(self.is_builtB, shape=[-1]),1e-10,1.0))+(1-self.target_built)*tf.log(tf.clip_by_value(1-tf.reshape(self.is_builtB, shape=[-1]),1e-10,1.0)))\n",
    "            self.policy_loss = - tf.reduce_sum(tf.log(tf.clip_by_value(self.responsible_outputs,1e-15,1.0)) * self.advantages)\n",
    "            self.valid_loss= - tf.reduce_sum(tf.log(tf.clip_by_value(self.validsB,1e-10,1.0)) * self.valids+tf.log(tf.clip_by_value(1-self.validsB,1e-10,1.0)) * (1-self.valids))\n",
    "            self.loss = 0.5 * self.value_loss + self.policy_loss + 0.5*self.block_loss + 0.5*self.built_loss + 0.5*self.valid_loss - self.entropy * 0.01\n",
    "\n",
    "            # Get gradients from local network using local losses and\n",
    "            # normalize the gradients using clipping\n",
    "            local_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope+'/qvaluesB')\n",
    "            self.gradients = tf.gradients(self.loss, local_vars)\n",
    "            self.var_norms = tf.global_norm(local_vars)\n",
    "            grads, self.grad_norms = tf.clip_by_global_norm(self.gradients, GRAD_CLIP)\n",
    "\n",
    "            # Apply local gradients to global network\n",
    "            global_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, GLOBAL_NET_SCOPE+'/qvaluesB')\n",
    "            self.apply_grads = trainer.apply_gradients(zip(grads, global_vars))\n",
    "\n",
    "            self.homogenize_weights = update_target_graph(str(scope)+'/qvaluesB', str(scope)+'/qvalues')\n",
    "        if TRAINING:\n",
    "            print(\"Hello World... From  \"+str(scope))     # :)\n",
    "\n",
    "    def _build_net(self,inputs):\n",
    "        w_init   = layers.variance_scaling_initializer()\n",
    "        conv1    = layers.conv3d(inputs=inputs, padding=\"SAME\", num_outputs=64, kernel_size=[4, 4, 4], stride=2, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=tf.nn.relu)\n",
    "        conv2    = layers.conv3d(inputs=conv1, padding=\"SAME\", num_outputs=128, kernel_size=[3, 2, 3], stride=1, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=tf.nn.relu)\n",
    "        conv3    = layers.conv3d(inputs=conv2, padding=\"SAME\", num_outputs=256, kernel_size=[3, 2, 3], stride=2, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=None)\n",
    "        res      = layers.conv3d(inputs=inputs, padding=\"VALID\", num_outputs=256, kernel_size=[4, 4, 4], stride=3, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=None)\n",
    "        conv4_in = tf.nn.relu(conv3+res)\n",
    "        conv4    = layers.conv3d(inputs=conv4_in, padding=\"SAME\", num_outputs=256, kernel_size=[2, 1, 2], stride=2, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=tf.nn.relu)\n",
    "        conv5    = layers.conv3d(inputs=conv4, padding=\"VALID\", num_outputs=512, kernel_size=[2, 1, 2], stride=1, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=None)\n",
    "        res2     = layers.conv3d(inputs=conv4_in, padding=\"VALID\", num_outputs=512, kernel_size=[3, 1, 3], stride=1, data_format=\"NDHWC\", weights_initializer=w_init,activation_fn=None)\n",
    "        flat     = tf.nn.relu(layers.flatten(conv5+res2))\n",
    "\n",
    "        h1 = layers.fully_connected(inputs=flat,  num_outputs=512)\n",
    "        h2 = layers.fully_connected(inputs=h1,  num_outputs=512, activation_fn=None)\n",
    "        h3 = tf.nn.relu(h2+flat)\n",
    "\n",
    "        #Recurrent network for temporal dependencies\n",
    "        lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(512,state_is_tuple=True)\n",
    "        c_init = np.zeros((1, lstm_cell.state_size.c), np.float32)\n",
    "        h_init = np.zeros((1, lstm_cell.state_size.h), np.float32)\n",
    "        state_init = [c_init, h_init]\n",
    "        c_in = tf.placeholder(tf.float32, [1, lstm_cell.state_size.c])\n",
    "        h_in = tf.placeholder(tf.float32, [1, lstm_cell.state_size.h])\n",
    "        state_in = (c_in, h_in)\n",
    "        rnn_in = tf.expand_dims(h3, [0])\n",
    "        step_size = tf.shape(inputs)[:1]\n",
    "        state_in = tf.nn.rnn_cell.LSTMStateTuple(c_in, h_in)\n",
    "        lstm_outputs, lstm_state = tf.nn.dynamic_rnn(\n",
    "        lstm_cell, rnn_in, initial_state=state_in, sequence_length=step_size,\n",
    "        time_major=False)\n",
    "        lstm_c, lstm_h = lstm_state\n",
    "        state_out = (lstm_c[:1, :], lstm_h[:1, :])\n",
    "        rnn_out = tf.reshape(lstm_outputs, [-1, 512])\n",
    "\n",
    "        policy_layer = layers.fully_connected(inputs=rnn_out, num_outputs=a_size,weights_initializer=normalized_columns_initializer(1./float(a_size)), biases_initializer=None, activation_fn=None)\n",
    "        policy=tf.nn.softmax(policy_layer)\n",
    "        policy_sig=tf.sigmoid(policy_layer)\n",
    "        value = layers.fully_connected(inputs=rnn_out, num_outputs=1,weights_initializer=normalized_columns_initializer(1.0), biases_initializer=None, activation_fn=None)\n",
    "        has_block = layers.fully_connected(inputs=rnn_out, num_outputs=1,weights_initializer=normalized_columns_initializer(1.0), biases_initializer=None, activation_fn=tf.sigmoid)\n",
    "        is_built  = layers.fully_connected(inputs=rnn_out, num_outputs=1,weights_initializer=normalized_columns_initializer(1.0), biases_initializer=None, activation_fn=tf.sigmoid)\n",
    "\n",
    "        return policy, value, state_out ,state_in, state_init, has_block, policy_sig, is_built"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Worker Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class Worker:\n",
    "    def __init__(self, game, metaAgentID, workerID, a_size, groupLock):\n",
    "        self.workerID = workerID\n",
    "        self.env = game\n",
    "        self.metaAgentID = metaAgentID\n",
    "        self.name = \"worker_\"+str(workerID)\n",
    "        self.agentID = ((workerID-1) % num_workers) + 1 \n",
    "        self.groupLock = groupLock\n",
    "        \n",
    "        self.nextGIF = episode_count # For GIFs output\n",
    "        #Create the local copy of the network and the tensorflow op to copy global parameters to local network\n",
    "        self.local_AC = ACNet(self.name,a_size,trainer)\n",
    "        self.copy_weights = self.local_AC.homogenize_weights\n",
    "        self.pull_global = update_target_graph(GLOBAL_NET_SCOPE, self.name)\n",
    "        \n",
    "    def train(self, rollout, sess, gamma, bootstrap_value):\n",
    "#       [s,a,r,s1,d,v[0,0],train_valid,pred_has_block,int(has_block),train_val,int(is_built)]\n",
    "        rollout = np.array(rollout)\n",
    "        observations = rollout[:,0]\n",
    "        actions = rollout[:,1]\n",
    "        rewards = rollout[:,2]\n",
    "        lastDone = rollout[-1,4]\n",
    "        values = rollout[:,5]\n",
    "        valids = rollout[:,6]\n",
    "        pred_has = rollout[:,7]\n",
    "        has_blocks = rollout[:,8]\n",
    "        train_value = rollout[:,9]\n",
    "        is_built = rollout[:,10]\n",
    "\n",
    "        # Here we take the rewards and values from the rollout, and use them to \n",
    "        # generate the advantage and discounted returns. (With bootstrapping)\n",
    "        # The advantage function uses \"Generalized Advantage Estimation\"\n",
    "        self.rewards_plus = np.asarray(rewards.tolist() + [bootstrap_value])\n",
    "        discounted_rewards = discount(self.rewards_plus,gamma)[:-1]\n",
    "        self.value_plus = np.asarray(values.tolist() + [bootstrap_value])\n",
    "        advantages = rewards + gamma * self.value_plus[1:] - self.value_plus[:-1]\n",
    "        advantages = good_discount(advantages,gamma)\n",
    "\n",
    "#         if not lastDone:\n",
    "        num_samples = min(EPISODE_SAMPLES,len(advantages))\n",
    "        sampleInd = np.sort(np.random.choice(advantages.shape[0], size=(num_samples,), replace=False))\n",
    "\n",
    "        # Update the global network using gradients from loss\n",
    "        # Generate network statistics to periodically save\n",
    "        rnn_state = self.local_AC.state_init\n",
    "        feed_dict = {self.local_AC.target_v:np.stack(discounted_rewards),\n",
    "            self.local_AC.inputsB:np.stack(observations),\n",
    "            self.local_AC.actions:actions,\n",
    "            self.local_AC.valids:np.stack(valids),\n",
    "            self.local_AC.advantages:advantages,\n",
    "            self.local_AC.train_value:train_value,\n",
    "            self.local_AC.has_blockB:np.reshape(pred_has,[np.shape(pred_has)[0],1]),\n",
    "            self.local_AC.target_has:has_blocks,\n",
    "            self.local_AC.target_built:is_built,\n",
    "            self.local_AC.state_inB[0]:rnn_state[0],\n",
    "            self.local_AC.state_inB[1]:rnn_state[1]}\n",
    "        \n",
    "        v_l,p_l,b_l,valid_l,e_l,g_n,v_n,bp_l,_ = sess.run([self.local_AC.value_loss,\n",
    "            self.local_AC.policy_loss,\n",
    "            self.local_AC.block_loss,\n",
    "            self.local_AC.valid_loss,\n",
    "            self.local_AC.entropy,\n",
    "            self.local_AC.grad_norms,\n",
    "            self.local_AC.var_norms,\n",
    "            self.local_AC.built_loss,\n",
    "            self.local_AC.apply_grads],\n",
    "            feed_dict=feed_dict)\n",
    "        return v_l / len(rollout), p_l / len(rollout),b_l / len(rollout),valid_l/len(rollout), e_l / len(rollout), g_n, v_n, bp_l\n",
    "\n",
    "    def shouldRun(self, coord, episode_count):\n",
    "        if TRAINING:\n",
    "            return (not coord.should_stop())\n",
    "        else:\n",
    "            return (episode_count < NUM_EXPS)\n",
    "\n",
    "    def work(self,max_episode_length,gamma,sess,coord,saver):\n",
    "        global episode_count, episode_rewards, episode_lengths, episode_mean_values, episode_invalid_ops\n",
    "        total_steps = 0\n",
    "\n",
    "        with sess.as_default(), sess.graph.as_default():\n",
    "            while self.shouldRun(coord, episode_count):\n",
    "                sess.run(self.pull_global)\n",
    "                sess.run(self.copy_weights)\n",
    "\n",
    "                episode_buffer, episode_values = [], []\n",
    "                episode_reward = episode_step_count = episode_inv_count = 0\n",
    "                if not TRAINING:\n",
    "                    completedQ[self.metaAgentID], d = False, False\n",
    "                    completed_time[self.metaAgentID] = np.nan\n",
    "                \n",
    "                # Initial state from the environment\n",
    "                if FULL_PLAN and random.random() < 0.5:\n",
    "                    validActions,has_block,is_built = self.env._reset(self.agentID, empty=EMPTY, full=True)\n",
    "                else:\n",
    "                    validActions,has_block,is_built = self.env._reset(self.agentID, empty=EMPTY, full=False)\n",
    "                s = self.env._observe(self.agentID)\n",
    "                rnn_state = self.local_AC.state_init\n",
    "                RewardNb = wrong_block = wrong_built = 0\n",
    "\n",
    "                saveGIF = False\n",
    "                if OUTPUT_GIFS and self.workerID == 1 and ((not TRAINING) or (episode_count >= self.nextGIF and episode_count >= OBSERVE_PERIOD)):\n",
    "                    saveGIF = True\n",
    "                    self.nextGIF += 128\n",
    "                    GIF_episode = int(episode_count)\n",
    "                    episode_frames = [ self.env._render(mode='rgb_array') ]\n",
    "\n",
    "                self.groupLock.release(0,self.name)\n",
    "                self.groupLock.acquire(1,self.name) # synchronize starting time of the threads\n",
    "\n",
    "                while (not self.env.finished): # Give me something!\n",
    "                    #Take an action using probabilities from policy network output.\n",
    "                    a_dist,v,rnn_state,pred_has_block,pred_is_built = sess.run([self.local_AC.policy,self.local_AC.value,self.local_AC.state_out,self.local_AC.has_block,self.local_AC.is_built], \n",
    "                                         feed_dict={self.local_AC.inputs:[s]\n",
    "                                                    ,self.local_AC.state_in[0]:rnn_state[0]\n",
    "                                                    ,self.local_AC.state_in[1]:rnn_state[1]})\n",
    "\n",
    "                    if(not (np.argmax(a_dist.flatten()) in validActions)):\n",
    "                        episode_inv_count+=1\n",
    "                    train_valid=np.zeros(a_size)\n",
    "                    train_valid[validActions]=1\n",
    "                    mysum=np.sum(train_valid)\n",
    "\n",
    "                    valid_dist = np.array([a_dist[0,validActions]])\n",
    "                    valid_dist /= np.sum(valid_dist)\n",
    "\n",
    "                    if TRAINING:\n",
    "                        if  (pred_has_block.flatten()[0] < 0.5) == has_block:\n",
    "                            wrong_block += 1\n",
    "                            a = np.random.choice(validActions)\n",
    "                            train_val = 0\n",
    "                        elif (pred_is_built.flatten()[0] < 0.5) == is_built:\n",
    "                            wrong_built += 1\n",
    "                            a = validActions[ np.random.choice(range(valid_dist.shape[1]),p=valid_dist.ravel()) ]\n",
    "                            train_val = 1.\n",
    "                        else:\n",
    "                            a = validActions[ np.random.choice(range(valid_dist.shape[1]),p=valid_dist.ravel()) ]\n",
    "                            train_val = 1.\n",
    "                    else:\n",
    "                        if GREEDY:\n",
    "                            a = np.argmax(a_dist.flatten())\n",
    "                        if not GREEDY or a not in validActions:\n",
    "                            a = validActions[ np.random.choice(range(valid_dist.shape[1]),p=valid_dist.ravel()) ]\n",
    "                        train_val = 1.\n",
    "\n",
    "                    _, r, d, validActions, has_block1, is_built1 = self.env.step((self.agentID, a))\n",
    "\n",
    "                    if not TRAINING:\n",
    "                        extraBlocks = max(0, self.env.world.countExtraBlocks(self.env.state_obj))\n",
    "                        if np.isnan(completed_time[self.metaAgentID]) and completedQ[self.metaAgentID] != is_built1 and is_built1:\n",
    "                            completed_time[self.metaAgentID] = episode_step_count+1\n",
    "                            scaffoldings[self.metaAgentID]   = extraBlocks\n",
    "                            blocks_left[self.metaAgentID]    = extraBlocks\n",
    "                        elif is_built1 and not np.isnan(blocks_left[self.metaAgentID]) and extraBlocks < blocks_left[self.metaAgentID]:\n",
    "                            blocks_left[self.metaAgentID]    = extraBlocks\n",
    "                        completedQ[self.metaAgentID]        |= is_built1\n",
    "\n",
    "                    self.groupLock.release(1,self.name)\n",
    "                    self.groupLock.acquire(0,self.name) # synchronize threads\n",
    "\n",
    "                    # Get common observation for all agents after all individual actions have been performed\n",
    "                    s1 = self.env._observe(self.agentID)\n",
    "                    d  = self.env.finished\n",
    "\n",
    "                    if saveGIF:\n",
    "                        episode_frames.append(self.env._render(mode='rgb_array'))\n",
    "\n",
    "                    self.groupLock.release(0,self.name)\n",
    "                    self.groupLock.acquire(1,self.name) # synchronize threads\n",
    "\n",
    "                    episode_buffer.append([s,a,r,s1,d,v[0,0],train_valid,pred_has_block,int(has_block),train_val,int(is_built)])\n",
    "                    episode_values.append(v[0,0])\n",
    "                    episode_reward += r\n",
    "                    s = s1\n",
    "                    total_steps += 1\n",
    "                    has_block = has_block1\n",
    "                    is_built = is_built1\n",
    "                    episode_step_count += 1\n",
    "\n",
    "                    if r>0:\n",
    "                        RewardNb += 1\n",
    "                    if d == True and TRAINING:\n",
    "                        print('\\n{} Goodbye World. We did it!'.format(episode_step_count), end='\\n')\n",
    "\n",
    "                    # If the episode hasn't ended, but the experience buffer is full, then we\n",
    "                    # make an update step using that experience rollout.\n",
    "                    if TRAINING and (len(episode_buffer) % EXPERIENCE_BUFFER_SIZE == 0 or d):\n",
    "                        # Since we don't know what the true final return is, we \"bootstrap\" from our current value estimation.\n",
    "                        if len(episode_buffer) >= EXPERIENCE_BUFFER_SIZE:\n",
    "\n",
    "                            if d:\n",
    "                                s1Value = 0.0\n",
    "                            else:\n",
    "                                s1Value = sess.run(self.local_AC.value, \n",
    "                                     feed_dict={self.local_AC.inputs:np.array([s])\n",
    "                                                ,self.local_AC.state_in[0]:rnn_state[0]\n",
    "                                                ,self.local_AC.state_in[1]:rnn_state[1]})[0,0]\n",
    "\n",
    "                            v_l,p_l,b_l,valid_l,e_l,g_n,v_n,bp_l = self.train(episode_buffer[-EXPERIENCE_BUFFER_SIZE:],sess,gamma,s1Value)\n",
    "\n",
    "                            sess.run(self.pull_global)\n",
    "                            sess.run(self.copy_weights)\n",
    "\n",
    "                    if episode_step_count >= max_episode_length or d:\n",
    "                        break\n",
    "\n",
    "                episode_rewards[self.metaAgentID].append(episode_reward)\n",
    "                episode_lengths[self.metaAgentID].append(episode_step_count)\n",
    "                episode_mean_values[self.metaAgentID].append(np.nanmean(episode_values))\n",
    "                episode_invalid_ops[self.metaAgentID].append(episode_inv_count)\n",
    "\n",
    "                # Periodically save gifs of episodes, model parameters, and summary statistics.\n",
    "                if episode_count % EXPERIENCE_BUFFER_SIZE == 0 and printQ:\n",
    "                    print('                                                                                   ', end='\\r')\n",
    "                    print('{} Episode terminated ({},{})'.format(episode_count, self.agentID, RewardNb), end='\\r')\n",
    "\n",
    "                if TRAINING:\n",
    "                    episode_count += 1. / num_workers\n",
    "\n",
    "                    if episode_count % SUMMARY_WINDOW == 0:\n",
    "                        if episode_count % 500 == 0:\n",
    "                            saver.save(sess, model_path+'/model-'+str(int(episode_count))+'.cptk')\n",
    "                            print ('Saved Model', end='\\r')\n",
    "                        mean_reward = np.mean(episode_rewards[self.metaAgentID][-SUMMARY_WINDOW:])\n",
    "                        mean_length = np.mean(episode_lengths[self.metaAgentID][-SUMMARY_WINDOW:])\n",
    "                        mean_value = np.mean(episode_mean_values[self.metaAgentID][-SUMMARY_WINDOW:])\n",
    "                        mean_invalid = np.mean(episode_invalid_ops[self.metaAgentID][-SUMMARY_WINDOW:])\n",
    "                        summary = tf.Summary()\n",
    "                        summary.value.add(tag='Perf/Reward', simple_value=mean_reward)\n",
    "                        summary.value.add(tag='Perf/Length', simple_value=mean_length)\n",
    "                        summary.value.add(tag='Perf/Valid Rate', simple_value=(mean_length-mean_invalid)/mean_length)\n",
    "                        summary.value.add(tag='Perf/Block Prediction Accuracy', simple_value=float(episode_step_count-wrong_block)/float(episode_step_count))\n",
    "                        summary.value.add(tag='Perf/Plan Completion Accuracy', simple_value=float(episode_step_count-wrong_built)/float(episode_step_count))\n",
    "\n",
    "                        summary.value.add(tag='Losses/Value Loss', simple_value=v_l)\n",
    "                        summary.value.add(tag='Losses/Policy Loss', simple_value=p_l)\n",
    "                        summary.value.add(tag='Losses/Plan Completion Loss', simple_value=bp_l)\n",
    "                        summary.value.add(tag='Losses/Block Prediction Loss', simple_value=b_l)\n",
    "                        summary.value.add(tag='Losses/Valid Loss', simple_value=valid_l)\n",
    "                        summary.value.add(tag='Losses/Grad Norm', simple_value=g_n)\n",
    "                        summary.value.add(tag='Losses/Var Norm', simple_value=v_n)\n",
    "                        global_summary.add_summary(summary, int(episode_count))\n",
    "\n",
    "                        global_summary.flush()\n",
    "\n",
    "                        if printQ:\n",
    "                            print('{} Tensorboard updated ({})'.format(episode_count, self.workerID), end='\\r')\n",
    "                elif not TRAINING and self.workerID == 1:\n",
    "                    if episode_buffer[0][-1] == 0: # only car about plan completion if init state didn't contain the completed plan...\n",
    "                        completed[episode_count]      = int(completedQ[self.metaAgentID])\n",
    "                    if not np.isnan(completed_time[self.metaAgentID]):\n",
    "                        plan_durations[episode_count] = completed_time[self.metaAgentID]\n",
    "                    len_episodes[episode_count]       = episode_step_count\n",
    "#                     saveGIF &= (episode_step_count < max_episode_length)\n",
    "\n",
    "                    episode_count += 1\n",
    "                    GIF_episode = int(episode_count)\n",
    "                    print('({}) Thread {}: {} steps ({} invalids).'.format(episode_count, self.workerID, episode_step_count, episode_inv_count))\n",
    "\n",
    "                self.groupLock.release(1,self.name)\n",
    "                self.groupLock.acquire(0,self.name) # synchronize threads\n",
    "\n",
    "                if saveGIF:\n",
    "                    # Dump episode frames for external gif generation (otherwise, makes the jupyter kernel crash)\n",
    "                    time_per_step = 0.1\n",
    "                    images = np.array(episode_frames)\n",
    "                    if TRAINING:\n",
    "                        gif_creation = lambda: make_gif(images, '{}/episode_{:d}_{:d}_{:.1f}.gif'.format(gifs_path,GIF_episode,episode_step_count,episode_reward), duration=len(images)*time_per_step,true_image=True,salience=False)\n",
    "                        threading.Thread(target=(gif_creation)).start()\n",
    "                    else:\n",
    "                        make_gif(images, '{}/episode_{:d}_{:d}.gif'.format(gifs_path,GIF_episode,episode_step_count), duration=len(images)*time_per_step,true_image=True,salience=False)\n",
    "                if self.workerID == 1 and SAVE_EPISODE_BUFFER and episode_step_count < max_episode_length:\n",
    "                    with open('{}/episode_{}.dat'.format(episodes_path,GIF_episode), 'wb') as file:\n",
    "                        pickle.dump(episode_buffer, file)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Learning parameters\n",
    "max_episode_length     = 1000\n",
    "episode_count          = 0\n",
    "EPISODE_START          = episode_count\n",
    "gamma                  = .9 # discount rate for advantage estimation and reward discounting\n",
    "GRAD_CLIP              = 300.0\n",
    "LR_Q                   = 2.e-5 # default: 2e-5\n",
    "a_size                 = 13 # New approach\n",
    "EXPERIENCE_BUFFER_SIZE = 100\n",
    "OBSERVE_PERIOD         = 0. # Period of pure observation (no value learning)\n",
    "SUMMARY_WINDOW         = 25\n",
    "NUM_META_AGENTS        = 4\n",
    "NUM_THREADS            = 4 #int(multiprocessing.cpu_count() / (2 * NUM_META_AGENTS))\n",
    "EPISODE_SAMPLES        = EXPERIENCE_BUFFER_SIZE\n",
    "load_model             = False\n",
    "RESET_TRAINER          = False\n",
    "model_path             = './model'\n",
    "gifs_path              = './gifs'\n",
    "train_path             = 'train_RNN'\n",
    "episodes_path          = 'gifs3D' # used to export episode_buffers that can be read/played/recorded by the visualizer\n",
    "GLOBAL_NET_SCOPE       = 'global'\n",
    "\n",
    "# Simulation options\n",
    "FULL_HELP              = False\n",
    "MAP_ID                 = 0 #0: RANDOMIZED_PLAN, other: given map (list in minecraft_SA4H.py)\n",
    "OUTPUT_GIFS            = False\n",
    "SAVE_EPISODE_BUFFER    = False\n",
    "\n",
    "# Testing\n",
    "TRAINING               = True\n",
    "GREEDY                 = False\n",
    "NUM_EXPS               = 20\n",
    "EMPTY                  = True and (not TRAINING)\n",
    "FULL_PLAN              = False # Help training cleanup by forcing all episodes\n",
    "MODEL_NUMBER           = 67000 # to start with the structure already completed\n",
    "                               # Should be enabled near the end of training\n",
    "\n",
    "# Shared arrays for tensorboard\n",
    "episode_rewards        = [ []     for _ in range(NUM_META_AGENTS) ]\n",
    "episode_lengths        = [ []     for _ in range(NUM_META_AGENTS) ]\n",
    "episode_mean_values    = [ []     for _ in range(NUM_META_AGENTS) ]\n",
    "episode_invalid_ops    = [ []     for _ in range(NUM_META_AGENTS) ]\n",
    "completedQ             = [ False  for _ in range(NUM_META_AGENTS) ]\n",
    "completed_time         = [ np.nan for _ in range(NUM_META_AGENTS) ]\n",
    "printQ                 = False # (for headless)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "print(\"Hello World\")\n",
    "if not os.path.exists(model_path):\n",
    "    os.makedirs(model_path)\n",
    "config = tf.ConfigProto(allow_soft_placement = True)\n",
    "config.gpu_options.allow_growth=True\n",
    "\n",
    "if not TRAINING:\n",
    "    blocks_left    = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "    scaffoldings   = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "    completed      = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "    plan_durations = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "    len_episodes   = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "    mutex = threading.Lock()\n",
    "    gifs_path += '_tests'\n",
    "    if SAVE_EPISODE_BUFFER and not os.path.exists(episodes_path):\n",
    "        os.makedirs(episodes_path)\n",
    "\n",
    "#Create a directory to save episode playback gifs to\n",
    "if OUTPUT_GIFS and not os.path.exists(gifs_path):\n",
    "    os.makedirs(gifs_path)\n",
    "\n",
    "with tf.device(\"/gpu:0\"):\n",
    "    master_network = ACNet(GLOBAL_NET_SCOPE,a_size,None) # Generate global network\n",
    "    trainer = tf.contrib.opt.NadamOptimizer(learning_rate=LR_Q, use_locking=True)\n",
    "\n",
    "    num_workers = NUM_THREADS # Set workers # = # of available CPU threads\n",
    "    if not TRAINING:\n",
    "        NUM_META_AGENTS = 1\n",
    "\n",
    "    gameEnvs, workers, groupLocks = [], [], []\n",
    "    for ma in range(NUM_META_AGENTS):\n",
    "        gameEnv = minecraft.MinecraftEnv(num_workers, observation_range=-1, observation_mode='default', FULL_HELP=FULL_HELP, MAP_ID=MAP_ID)\n",
    "        gameEnvs.append(gameEnv)\n",
    "\n",
    "        # Create groupLock\n",
    "        workerNames = [\"worker_\"+str(i) for i in range(ma*num_workers+1,(ma+1)*num_workers+1)]\n",
    "        groupLock = GroupLock.GroupLock([workerNames,workerNames])\n",
    "        groupLocks.append(groupLock)\n",
    "\n",
    "        # Create worker classes\n",
    "        workersTmp = []\n",
    "        for i in range(ma*num_workers+1,(ma+1)*num_workers+1):\n",
    "            workersTmp.append(Worker(gameEnv,ma,i,a_size,groupLock))\n",
    "        workers.append(workersTmp)\n",
    "\n",
    "    if TRAINING:\n",
    "        global_summary = tf.summary.FileWriter(train_path)\n",
    "    else:\n",
    "        global_summary = 0\n",
    "    saver = tf.train.Saver(max_to_keep=5)\n",
    "\n",
    "    with tf.Session(config=config) as sess:\n",
    "        coord = tf.train.Coordinator()\n",
    "        if load_model == True:\n",
    "            print ('Loading Model...')\n",
    "            if not TRAINING:\n",
    "                with open(model_path+'/checkpoint', 'w') as file:\n",
    "                    file.write('model_checkpoint_path: \"model-{}.cptk\"'.format(MODEL_NUMBER))\n",
    "                    file.close()\n",
    "            ckpt = tf.train.get_checkpoint_state(model_path)\n",
    "            saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "            if RESET_TRAINER:\n",
    "                trainer = tf.contrib.opt.NadamOptimizer(learning_rate=LR_Q, use_locking=True)\n",
    "        else:\n",
    "            sess.run(tf.global_variables_initializer())\n",
    "\n",
    "        # This is where the asynchronous magic happens.\n",
    "        # Start the \"work\" process for each worker in a separate thread.\n",
    "        worker_threads = []\n",
    "        for ma in range(NUM_META_AGENTS):\n",
    "            for worker in workers[ma]:\n",
    "                groupLocks[ma].acquire(0,worker.name) # synchronize starting time of the threads\n",
    "                worker_work = lambda: worker.work(max_episode_length,gamma,sess,coord,saver)\n",
    "                print(\"Starting worker \" + str(worker.workerID))\n",
    "                t = threading.Thread(target=(worker_work))\n",
    "                t.start()\n",
    "                worker_threads.append(t)\n",
    "        coord.join(worker_threads)\n",
    "\n",
    "if not TRAINING:\n",
    "    print('[{:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}]'.format(\n",
    "           1 - np.nanmean(blocks_left/scaffoldings), np.sqrt(np.nanvar(blocks_left/scaffoldings)),\n",
    "           np.nanmean(completed), np.sqrt(np.nanvar(completed)),\n",
    "           np.nanmean(plan_durations), np.sqrt(np.nanvar(plan_durations)),\n",
    "           np.nanmean(len_episodes), np.sqrt(np.nanvar(len_episodes)),\n",
    "           np.nanmean(np.asarray(len_episodes < max_episode_length, dtype=float)),\n",
    "           np.nanmean(blocks_left), np.sqrt(np.nanvar(blocks_left)),\n",
    "           np.nanmean(scaffoldings), np.sqrt(np.nanvar(scaffoldings)))\n",
    "         )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Systematic Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "GREEDY                 = False\n",
    "NUM_EXPS               = 100\n",
    "EMPTY                  = True\n",
    "\n",
    "SAVE_EPISODE_BUFFER    = True\n",
    "NUM_META_AGENTS        = 1\n",
    "\n",
    "if SAVE_EPISODE_BUFFER and not os.path.exists('gifs3D'):\n",
    "    os.makedirs('gifs3D')\n",
    "\n",
    "for num_workers in [1,2,4,8,12,16,3,6,10,14]:\n",
    "    for MAP_ID in range(1,7):\n",
    "        episode_count = 0\n",
    "\n",
    "        tf.reset_default_graph()\n",
    "        config = tf.ConfigProto(allow_soft_placement = True)\n",
    "        config.gpu_options.allow_growth=True\n",
    "\n",
    "        blocks_left    = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "        scaffoldings   = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "        completed      = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "        plan_durations = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "        len_episodes   = np.array([np.nan for _ in range(NUM_EXPS)])\n",
    "        mutex = threading.Lock()\n",
    "\n",
    "        episodes_path = 'gifs3D/{:d}_{:d}'.format(num_workers,MAP_ID)\n",
    "        if SAVE_EPISODE_BUFFER and not os.path.exists(episodes_path):\n",
    "            os.makedirs(episodes_path)\n",
    "\n",
    "        with tf.device(\"/gpu:0\"):\n",
    "            master_network = ACNet(GLOBAL_NET_SCOPE,a_size,None) # Generate global network\n",
    "            trainer = tf.contrib.opt.NadamOptimizer(learning_rate=LR_Q, use_locking=True)\n",
    "\n",
    "            global_summary = 0\n",
    "            saver = tf.train.Saver(max_to_keep=5)\n",
    "\n",
    "            with tf.Session(config=config) as sess:\n",
    "                coord = tf.train.Coordinator()\n",
    "                with open(model_path+'/checkpoint', 'w') as file:\n",
    "                    file.write('model_checkpoint_path: \"model-{}.cptk\"'.format(MODEL_NUMBER))\n",
    "                    file.close()\n",
    "                ckpt = tf.train.get_checkpoint_state(model_path)\n",
    "                saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "\n",
    "                gameEnvs, workers, groupLocks = [], [], []\n",
    "                for ma in range(NUM_META_AGENTS):\n",
    "                    gameEnv = minecraft.MinecraftEnv(num_workers, observation_range=-1, observation_mode='default', FULL_HELP=FULL_HELP, MAP_ID=MAP_ID)\n",
    "                    gameEnvs.append(gameEnv)\n",
    "\n",
    "                    # Create groupLock\n",
    "                    workerNames = [\"worker_\"+str(i) for i in range(ma*num_workers+1,(ma+1)*num_workers+1)]\n",
    "                    groupLock = GroupLock.GroupLock([workerNames,workerNames])\n",
    "                    groupLocks.append(groupLock)\n",
    "\n",
    "                    # Create worker classes\n",
    "                    workersTmp = []\n",
    "                    for i in range(ma*num_workers+1,(ma+1)*num_workers+1):\n",
    "                        workersTmp.append(Worker(gameEnv,ma,i,a_size,groupLock))\n",
    "                    workers.append(workersTmp)\n",
    "\n",
    "                # This is where the asynchronous magic happens.\n",
    "                # Start the \"work\" process for each worker in a separate thread.\n",
    "                worker_threads = []\n",
    "                for ma in range(NUM_META_AGENTS):\n",
    "                    for worker in workers[ma]:\n",
    "                        groupLocks[ma].acquire(0,worker.name) # synchronize starting time of the threads\n",
    "                        worker_work = lambda: worker.work(max_episode_length,gamma,sess,coord,saver)\n",
    "                        t = threading.Thread(target=(worker_work))\n",
    "                        t.start()\n",
    "                        worker_threads.append(t)\n",
    "                coord.join(worker_threads)\n",
    "\n",
    "        print('[{:d}, {:d}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}]'.format(\n",
    "               num_workers, MAP_ID,\n",
    "               1 - np.nanmean(blocks_left/scaffoldings), np.sqrt(np.nanvar(blocks_left/scaffoldings)),\n",
    "               np.nanmean(completed), np.sqrt(np.nanvar(completed)),\n",
    "               np.nanmean(plan_durations), np.sqrt(np.nanvar(plan_durations)),\n",
    "               np.nanmean(len_episodes), np.sqrt(np.nanvar(len_episodes)),\n",
    "               np.nanmean(np.asarray(len_episodes < max_episode_length, dtype=float)),\n",
    "               np.nanmean(blocks_left), np.sqrt(np.nanvar(blocks_left)),\n",
    "               np.nanmean(scaffoldings), np.sqrt(np.nanvar(scaffoldings)))\n",
    "             )\n",
    "\n",
    "        ofp = open('results.txt','a')\n",
    "        ofp.write('{:d}, {:d}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}\\n'.format(\n",
    "               num_workers, MAP_ID,\n",
    "               1 - np.nanmean(blocks_left/scaffoldings), np.sqrt(np.nanvar(blocks_left/scaffoldings)),\n",
    "               np.nanmean(completed), np.sqrt(np.nanvar(completed)),\n",
    "               np.nanmean(plan_durations), np.sqrt(np.nanvar(plan_durations)),\n",
    "               np.nanmean(len_episodes), np.sqrt(np.nanvar(len_episodes)),\n",
    "               np.nanmean(np.asarray(len_episodes < max_episode_length, dtype=float)),\n",
    "               np.nanmean(blocks_left), np.sqrt(np.nanvar(blocks_left)),\n",
    "               np.nanmean(scaffoldings), np.sqrt(np.nanvar(scaffoldings)))\n",
    "             )\n",
    "        ofp.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
